"""
https://www.kaggle.com/sachink1729/audio-classification-using-mfccs-as-features#5.-Create-X_train,-Y_train-&-X_test,-Y_test
"""

import pandas as pd
import numpy as np
import os
import librosa
import tensorflow as tf
import librosa.display
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
import tensorflow.keras.layers as layers
from sklearn.metrics import confusion_matrix, classification_report, accuracy_score, f1_score


def sep(label = '', cnt=32):
    print('-' * cnt, label, '-' * cnt, sep='')


sep('Arguments')
# assign directory
# directory = '../input/audio-cats-and-dogs/cats_dogs/train/'
dir = '../../../../large_data/audio/_many_files/cat_dog_archive/cats_dogs/train/'
dir_train = dir
# assign directory
# directory = '../input/audio-cats-and-dogs/cats_dogs/test/'
dir = '../../../../large_data/audio/_many_files/cat_dog_archive/cats_dogs/test/'
dir_test = dir

# train data creator
sep('train data creator')
directory = dir_train

ID = []
label = []

# iterate over files in
# that directory

for folder in os.listdir(directory):  # go into the directory
    for filename in os.listdir(directory + str(folder)):  # go in every class
        # f = os.path.join(directory + str(folder), filename)  # scan through every file in that class
        f = directory + str(folder) + '/' + filename
        if os.path.isfile(f):
            ID.append(f.split('/')[-1])
            label.append(f.split('/')[-2])
train_data = pd.DataFrame()
train_data['ID'] = ID
train_data['label'] = label

print(train_data)

# plt.figure(figsize=(8, 6), dpi=80)
# sns.set_theme(style="darkgrid")
# sns.countplot('label', data=train_data)
# plt.title('counts: \n' + 'dogs:' + str(train_data.label.value_counts()[1]) +
#           '\n cats:' + str(train_data.label.value_counts()[0]))
# plt.show()

# create test data
sep('create test data')
directory = dir_test

ID = []
label = []

# iterate over files in
# that directory

for folder in os.listdir(directory):
    for filename in os.listdir(directory + str(folder)):
        # f = os.path.join(directory + str(folder), filename)
        f = directory + str(folder) + '/' + filename
        if os.path.isfile(f):
            ID.append(f.split('/')[-1])
            label.append(f.split('/')[-2])

for i in range(len(label)):
    if (label[i] == 'test'):
        label[i] = 'dogs'

test_data = pd.DataFrame()
test_data['ID'] = ID
test_data['label'] = label

print(test_data['label'].value_counts())
print(test_data)

sep('Feature extraction of audio files using MFCCs')


def extract_features(directory):
    features = []
    ID = []
    for folder in os.listdir(directory):
        for filename in os.listdir(directory + str(folder)):
            # f = os.path.join(directory + str(folder), filename)
            f = directory + str(folder) + '/' + filename
            if os.path.isfile(f):
                x, sr = librosa.load(f, res_type='kaiser_fast', sr=None)
                mfccs = np.mean(librosa.feature.mfcc(x, sr=sr, n_mfcc=100).T, axis=0)
                features.append(mfccs)
                ID.append(f.split('/')[-1])
    return [ID, features]


# ID, features_train=extract_features('../input/audio-cats-and-dogs/cats_dogs/train/')
ID_train, features_train = extract_features(dir_train)

# ID,features_test=extract_features('../input/audio-cats-and-dogs/cats_dogs/test/')
ID_test, features_test = extract_features(dir_test)

sep('Create X_train, Y_train & X_test, Y_test')
X_train = np.array(features_train)
X_test = np.array(features_test)

Y_train = train_data.label
Y_test = test_data.label

le = LabelEncoder()

temp = le.fit_transform(Y_train)
Y_train = temp.reshape(-1, 1)

temp = le.fit_transform(Y_test)
Y_test = temp.reshape(-1, 1)
print(X_train.shape, X_test.shape)
print(Y_train.shape, Y_test.shape)

sep('Build audio classification model')

model = tf.keras.Sequential()
model.add(layers.Dense(input_shape=(100,), units=200, activation='relu'))
model.add(layers.Dense(200, activation='relu'))
model.add(layers.Dense(200, activation='relu'))
model.add(layers.Dense(1, activation='sigmoid'))

model.summary()

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

hist = model.fit(X_train, Y_train, epochs=100)

spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Train Loss values over 100 epochs:')
plt.plot(hist.history['loss'], color='red', linewidth=2)

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Train accuracy values over 100 epochs:')
plt.plot(hist.history['accuracy'], color='purple')

plt.show()

sep('Testing the model with test data and metrics')

Y_pred = model.predict(X_test)
# since we are using sigmoid activation function at the output layer
Y_pred = (Y_pred > 0.5) * 1

print(classification_report(Y_test, Y_pred))

plt.figure(figsize=(10, 8), dpi=80)
sns.heatmap(confusion_matrix(Y_test, Y_pred), annot=True, cmap='Blues')
plt.title('1 signifies dog sounds and 0 signifies cat sounds \n' + 'Accuracy:' + str(accuracy_score(Y_test, Y_pred)))
plt.show()
